# encoding = utf8
import re
import math
import codecs
import random

import numpy as np
import jieba
import jieba.posseg as pseg
jieba.load_userdict("dictionary.txt") #add UMLS dictionary
jieba.initialize()


def create_dico(item_list):
    """
    Create a dictionary of items from a list of list of items.
    """
    assert type(item_list) is list
    dico = {}
    for items in item_list:
        for item in items:
            if item not in dico:
                dico[item] = 1
            else:
                dico[item] += 1
    return dico


def create_mapping(dico):
    """
    Create a mapping (item to ID / ID to item) from a dictionary.
    Items are ordered by decreasing frequency.
    """
    sorted_items = sorted(dico.items(), key=lambda x: (-x[1], x[0]))
    id_to_item = {i: v[0] for i, v in enumerate(sorted_items)}
    item_to_id = {v: k for k, v in id_to_item.items()}
    return item_to_id, id_to_item


def zero_digits(s):
    """
    Replace every digit in a string by a zero.
    """
    return re.sub('\d', '0', s)


def iob2(tags):
    """
    Check that tags have a valid IOB format.
    Tags in IOB1 format are converted to IOB2.
    """
    for i, tag in enumerate(tags):
        if tag == 'O':
            continue
        split = tag.split('-')
        if len(split) != 2 or split[0] not in ['I', 'B']:
            return False
        if split[0] == 'B':
            continue
        elif i == 0 or tags[i - 1] == 'O':  # conversion IOB1 to IOB2
            tags[i] = 'B' + tag[1:]
        elif tags[i - 1][1:] == tag[1:]:
            continue
        else:  # conversion IOB1 to IOB2
            tags[i] = 'B' + tag[1:]
    return True


def iob_iobes(tags):
    """
    IOB -> IOBES
    """
    new_tags = []
    for i, tag in enumerate(tags):
        if tag == 'O':
            new_tags.append(tag)
        elif tag.split('-')[0] == 'B':
            if i + 1 != len(tags) and \
               tags[i + 1].split('-')[0] == 'I':
                new_tags.append(tag)
            else:
                new_tags.append(tag.replace('B-', 'S-'))
        elif tag.split('-')[0] == 'I':
            if i + 1 < len(tags) and \
                    tags[i + 1].split('-')[0] == 'I':
                new_tags.append(tag)
            else:
                new_tags.append(tag.replace('I-', 'E-'))
        else:
            raise Exception('Invalid IOB format!')
    return new_tags


def iobes_iob(tags):
    """
    IOBES -> IOB
    """
    new_tags = []
    for i, tag in enumerate(tags):
        if tag.split('-')[0] == 'B':
            new_tags.append(tag)
        elif tag.split('-')[0] == 'I':
            new_tags.append(tag)
        elif tag.split('-')[0] == 'S':
            new_tags.append(tag.replace('S-', 'B-'))
        elif tag.split('-')[0] == 'E':
            new_tags.append(tag.replace('E-', 'I-'))
        elif tag.split('-')[0] == 'O':
            new_tags.append(tag)
        else:
            raise Exception('Invalid format!')
    return new_tags


def insert_singletons(words, singletons, p=0.5):
    """
    Replace singletons by the unknown word with a probability p.
    """
    new_words = []
    for word in words:
        if word in singletons and np.random.uniform() < p:
            new_words.append(0)
        else:
            new_words.append(word)
    return new_words


def get_seg_pos_features(string, pos_list):
    """
    Segment text with jieba
    """
    seg_feature = []
    pos_feature = []
    #jieba.load_userdict("dictionary.txt") #add UMLS dictionary
    words = pseg.cut(string)
    for word, flag in words:
        if len(word) == 1:
            seg_feature.append(0)
            pos_feature.append(pos_list[flag if flag in pos_list else 'un'])   
        else:
            tmp = [2] * len(word)
            tmp[0] = 1
            tmp[-1] = 3
            seg_feature.extend(tmp)
            tmp2=[pos_list[flag if flag in pos_list else 'un']] * len(word)
            pos_feature.extend(tmp2)
    #print(len(string))
    #print(len(seg_feature))    
    return seg_feature, pos_feature

#'''
def get_pos_list():
	tags=open('data/pos_list','r').readlines()
	pos_list={}
	i=0
	for t in tags:
		pos_list[t.strip()]=i
		i+=1
	return pos_list
#'''


#get the word id of each char in the string 
def get_words(string, word_to_id):
	words=[]
	#print(word_to_id)
        #jieba.load_userdict("dictionary.txt") #add UMLS dictionary
	for word in jieba.cut(string):
		#word=word if word in word_to_id else '<UNK>'
		if len(word) ==1:
			words.append(word_to_id[word if word in word_to_id else '<UNK>'])
		else:
			tmp=[word_to_id[word if word in word_to_id else '<UNK>']]*len(word)
			words.extend(tmp)
	return words


#update original embedding matrix with in the context of n window
def window_context(emb, windows):
	padding=windows//2
	emb=np.array(emb)
	lens=len(emb)
	print(emb)
	print (lens)
	for i in range(lens):
		if (i>=padding) and (i<(lens+padding)):
			emb[i]=np.sum(emb[i-padding : i+padding+1], axis=0)/windo
		elif i < padding :
			emb[i]=np.sum(emb[0 : i+padding+1], axis=0)/windows
		elif i > (lens+padding) :
			emb[i]=np.sum(emb[i-padding : lens], axis=0)/windows
		print("XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX")
		print(emb[i])	
	return emb


def contextwin(l, win):
    '''
    win :: int corresponding to the size of the window
    given a list of indexes composing a sentence

    l :: array containing the word indexes

    it will return a list of list of indexes corresponding
    to context windows surrounding each word in the sentence
    '''
    assert (win % 2) == 1
    assert win >= 1
    l = list(l)

    lpadded = win // 2 * [-1] + l + win // 2 * [-1]
    out = [lpadded[i:(i + win)] for i in range(len(l))]

    assert len(out) == len(l)
    return out


def create_input(data):
    """
    Take sentence data and return an input for
    the training or the evaluation function.
    """
    inputs = list()
    inputs.append(data['chars'])
    inputs.append(data['words'])
    inputs.append(data["segs"])
    inputs.append(data["poss"])
    inputs.append(data['tags'])
    return inputs

def create_segvector():
    seg_vector=[0]*4



#emb_weights_char = load_vec(config["emb_file"], id_to_char, config["char_dim"], emb_weights_char)
#def load_word2vec(emb_path, id_to_char, char_dim, old_weights):
def load_word2vec(emb_path, emb_path2, id_to_char, char_dim, old_weights):

    """
    Load char embedding from pre-trained file
    embedding size must match
    """
    new_weights = old_weights
    print('Loading pretrained char embeddings from {}...'.format(emb_path))
    pre_trained = {} #char emb file
    #pre_trained2 = {} #word emb file
    emb_invalid = 0

    #get char emb list
    for i, line in enumerate(codecs.open(emb_path, 'r', 'utf-8')):
        line = line.rstrip().split()
        if len(line) == char_dim + 1:
            pre_trained[line[0]] = np.array(
                [float(x) for x in line[1:]]
            ).astype(np.float32)
        else:
            emb_invalid += 1
    if emb_invalid > 0:
        print('WARNING: %i invalid lines in char vectors' % emb_invalid)

    #update pre_trained according to word emb list where word length=1
    '''
    update_num=0
    for i, line in enumerate(codecs.open(emb_path2, 'r', 'utf-8')):
        line = line.rstrip().split()
        w=line[0]
        if len(w)==1 and (w in pre_trained.keys()) and len(line) == char_dim + 1:
            wordemb=np.array([float(x) for x in line[1:]]).astype(np.float32)
            #print ('char: ' + w)
            #print (wordemb)            
            #print (pre_trained[w])
            pre_trained[w] = 0.5 * (pre_trained[w] + wordemb)
            update_num +=1
            #print (pre_trained[w])
    if update_num > 0:
        print('%i of %i char embs updated according to word emb list' % (update_num, len(pre_trained)))
    '''
    c_found = 0
    c_lower = 0
    c_zeros = 0
    n_words = len(id_to_char)
    #print(n_words)
    #print(len(new_weights))
    # Lookup table initialization
    for i in range(n_words):
        word = id_to_char[i]
        if word in pre_trained:
            new_weights[i] = pre_trained[word]
            c_found += 1
        elif word.lower() in pre_trained:
            new_weights[i] = pre_trained[word.lower()]
            c_lower += 1
        elif re.sub('\d', '0', word.lower()) in pre_trained:
            new_weights[i] = pre_trained[
                re.sub('\d', '0', word.lower())
            ]
            c_zeros += 1
    print('Loaded %i pretrained char embeddings.' % len(pre_trained))
    print('%i / %i (%.4f%%) chars have been initialized with '
          'pretrained embeddings.' % (
        c_found + c_lower + c_zeros, n_words,
        100. * (c_found + c_lower + c_zeros) / n_words)
    )
    print('%i found directly, %i after lowercasing, '
          '%i after lowercasing + zero.' % (
        c_found, c_lower, c_zeros
    ))
    return new_weights


#emb_weights_word = load_vec2(config["emb_file2"], id_to_word, config["word_dim"], emb_weights_word)
def load_word2vec2(emb_path, id_to_word, word_dim, old_weights):
    """
    Load word embedding from pre-trained file
    embedding size must match
    """
    new_weights = old_weights
    
    print('Loading pretrained word embeddings from {}...'.format(emb_path))
    pre_trained = {}
    emb_invalid = 0
    for i, line in enumerate(codecs.open(emb_path, 'r', 'utf-8')):
        line = line.rstrip().split()
        if len(line) == word_dim + 1:
            pre_trained[line[0]] = np.array(
                [float(x) for x in line[1:]]
            ).astype(np.float32)
        else:
            emb_invalid += 1
    if emb_invalid > 0:
        print('WARNING: %i invalid lines in word vectors' % emb_invalid)
    c_found = 0
    c_lower = 0
    c_zeros = 0
    n_words = len(id_to_word)
    #print(n_words)
    #print(len(new_weights))
    #print(new_weights[0])
    #print('#################')
    #print(new_weights[0][2])
    #print()
    #new_weights=np.zeros((n_words,seg_dim))
    #print(new_weights)
    # Lookup table initialization
    for i in range(n_words):
        word = id_to_word[i]
        if word in pre_trained:
            new_weights[i] = pre_trained[word]
            c_found += 1
        elif word.lower() in pre_trained:
            new_weights[i] = pre_trained[word.lower()]
            c_lower += 1
        elif re.sub('\d', '0', word.lower()) in pre_trained:
            new_weights[i] = pre_trained[
                re.sub('\d', '0', word.lower())
            ]
            c_zeros += 1
    print('Loaded %i pretrained word embeddings.' % len(pre_trained))
    print('%i / %i (%.4f%%) words have been initialized with '
          'pretrained embeddings.' % (
        c_found + c_lower + c_zeros, n_words,
        100. * (c_found + c_lower + c_zeros) / n_words)
    )
    print('%i found directly, %i after lowercasing, '
          '%i after lowercasing + zero.' % (
        c_found, c_lower, c_zeros
    ))
    return new_weights


#update seg_vec
def load_word2vec3(old_weights):
    """
    Load word embedding from pre-trained file
    embedding size must match
    """
    #print(old_weights)
    new_weights = old_weights
    
    #print('Loading pretrained word embeddings from {}...'.format(emb_path))
    print('Loading seg embeddings ...')
    new_weights[0] = [1, 0, 0, 0]
    new_weights[1] = [0, 1, 0, 0]
    new_weights[2] = [0, 0, 1, 0]
    new_weights[3] = [0, 0, 0, 1]
    print('seg embeddings loaded')
    print(new_weights)
    return new_weights


#update pos_vec
def load_word2vec4(old_weights):
    """
    Load word embedding from pre-trained file
    embedding size must match
    """
    #print('XXXXXXXXXX')
    #print(old_weights)
    new_weights = old_weights
    
    #print('Loading pretrained word embeddings from {}...'.format(emb_path))
    print('Loading pos embeddings ...')
    for i in range(40):
    	new_weights[i] = [0] * 40
    	new_weights[i][i]=1
    #print(new_weights)
    print('pos embeddings loaded')
    #print(new_weights)
    return new_weights


def full_to_half(s):
    """
    Convert full-width character to half-width one 
    """
    n = []
    for char in s:
        num = ord(char)
        if num == 0x3000:
            num = 32
        elif 0xFF01 <= num <= 0xFF5E:
            num -= 0xfee0
        char = chr(num)
        n.append(char)
    return ''.join(n)


def cut_to_sentence(text):
    """
    Cut text to sentences 
    """
    sentence = []
    sentences = []
    len_p = len(text)
    pre_cut = False
    for idx, word in enumerate(text):
        sentence.append(word)
        cut = False
        if pre_cut:
            cut=True
            pre_cut=False
        if word in u"。;!?\n":
            cut = True
            if len_p > idx+1:
                if text[idx+1] in ".。”\"\'“”‘’?!":
                    cut = False
                    pre_cut=True

        if cut:
            sentences.append(sentence)
            sentence = []
    if sentence:
        sentences.append("".join(list(sentence)))
    return sentences


def replace_html(s):
    s = s.replace('&quot;','"')
    s = s.replace('&amp;','&')
    s = s.replace('&lt;','<')
    s = s.replace('&gt;','>')
    s = s.replace('&nbsp;',' ')
    s = s.replace("&ldquo;", "“")
    s = s.replace("&rdquo;", "”")
    s = s.replace("&mdash;","")
    s = s.replace("\xa0", " ")
    return(s)


def input_from_line(line, char_to_id):
    """
    Take sentence data and return an input for
    the training or the evaluation function.
    """
    line = full_to_half(line)
    line = replace_html(line)
    inputs = list()
    inputs.append([line])
    line.replace(" ", "$")
    inputs.append([[char_to_id[char] if char in char_to_id else char_to_id["<UNK>"]
                   for char in line]])
    inputs.append([get_seg_features(line)])
    inputs.append([[]])
    return inputs


class BatchManager(object):

    def __init__(self, data,  batch_size):
        #print(len(data[0]))
        self.batch_data = self.sort_and_pad(data, batch_size)
        self.len_data = len(self.batch_data)
        #print(len(self.batch_data[0]))

    def sort_and_pad(self, data, batch_size):
        num_batch = int(math.ceil(len(data) /batch_size))
        sorted_data = sorted(data, key=lambda x: len(x[0]))
        batch_data = list()
        for i in range(num_batch):
            batch_data.append(self.pad_data(sorted_data[i*batch_size : (i+1)*batch_size]))
        return batch_data

    @staticmethod
    def pad_data(data):
        strings = []
        chars = []
        words = []
        segs = []
        poss = []
        targets = []
        max_length = max([len(sentence[0]) for sentence in data])
        for line in data:
            #print(line)
            string, char, word, seg, pos, target = line
            padding = [0] * (max_length - len(string))
            strings.append(string + padding)
            chars.append(char + padding)
            words.append(word + padding)
            segs.append(seg + padding)
            poss.append(pos + padding)
            targets.append(target + padding)
            #print("string len: %i / chars len: %i / words len: %i / segs len: %i / tags len: %i." % (
        #len(strings), len(chars), len(words), len(segs), len(targets)))
        return [strings, chars, words, segs, poss, targets]

    def iter_batch(self, shuffle=False):
        if shuffle:
            random.shuffle(self.batch_data)
        for idx in range(self.len_data):
            yield self.batch_data[idx]
